#ifndef BOOTSTRAPPING_KARATSUBA_H
#define BOOTSTRAPPING_KARATSUBA_H
#include <vector>
#include <cassert>
#include <memory>

/*
 * 我们使用karatsuba算法进行明文多项式的乘法
 */
namespace karatsuba {
inline size_t LeadingZeros(uint64_t mask)
{
    unsigned long result = 64;
    while (mask != 0) {
        mask = mask >> 1;
        result -= 1;
    }
    return result;
}

template <typename T, typename F> void ElementwiseAssign(T *dst, const T *src, size_t n, const F &f);

template <bool addAssign, typename T, typename Add, typename Sub, typename Mul>
void NaiveAssignMul(T *dst, const T *lhs, const T *rhs, size_t inSizeLog2, const Add &add, const Sub &sub,
    const Mul &mul);

template <unsigned inSizeLog2, unsigned thresholdLog2, bool addAssign, typename T, typename Add, typename Sub,
    typename Mul>
void KaratsubaAssignMulImpl(T *dst, const T *lhs, const T *rhs, T *memory, const Add &add, const Sub &sub,
    const Mul &mul);

template <unsigned thresholdLog2, bool addAssign, typename T, typename Add, typename Sub, typename Mul>
void DispatchKaratsubaAssignMulImpl(unsigned inSizeLog2, T *dst, const T *lhs, const T *rhs, T *memory, const Add &add,
    const Sub &sub, const Mul &mul);

template <unsigned thresholdLog2, typename T, typename Add, typename Sub, typename Mul>
void Karatsuba(T *dst, size_t dstLen, const T *lhs, size_t lhsLen, const T *rhs, size_t rhsLen, const Add &add,
    const Sub &sub, const Mul &mul);

template <typename T> size_t KaratsubaMemElementCount(size_t inSizeLog2, size_t thresholdLog2);

/*
 * @brief 执行逐元素赋值操作
 * 该函数对两个数组进行逐元素的赋值操作，将源数组 `src` 中的每个元素应用函数 `f` 后赋值给目标数组 `dst`
 * 中对应位置的元素。
 * @tparam T 元素的类型。
 * @tparam F 操作的类型，通常是一个函数或仿函数。
 *
 * @param dst 目标数组，存储操作后的结果。
 * @param src 源数组，将通过函数 `f` 中的操作提供输入。
 * @param n 数组的大小，即元素的数量。
 * @param f 逐元素操作的函数或仿函数，该函数接受两个 `T` 类型的参数，并返回一个 `T` 类型的结果。
 *
 * @note 该函数对两个数组的对应元素进行操作，`f` 会对目标数组中的每个元素与源数组中的相应元素进行处理。
 */
template <typename T, typename F> inline void ElementwiseAssign(T *dst, const T *src, size_t n, const F &f)
{
    for (size_t i = 0; i < n; ++i) {
        dst[i] = f(dst[i], src[i]);
    }
}

template <bool addAssign, typename T, typename Add, typename Sub, typename Mul>
inline void NaiveAssignMul(T *dst, const T *lhs, const T *rhs, size_t inSizeLog2, const Add &add, const Sub &sub,
    const Mul &mul)
{
    const size_t n = (size_t)1 << inSizeLog2;
    for (size_t i = 0; i < 2 * n; ++i) {
        T acc = 0;
        const size_t start =
            static_cast<size_t>(std::max<int64_t>(static_cast<int64_t>(i) - static_cast<int64_t>(n) + 1, 0));
        for (size_t j = start; j < std::min(n, i + 1); ++j) {
            acc = add(acc, mul(lhs[i - j], rhs[j]));
        }
        if constexpr (addAssign) {
            dst[i] = add(dst[i], acc);
        } else {
            dst[i] = acc;
        }
    }
}

template <unsigned inSizeLog2, unsigned thresholdLog2, bool addAssign, typename T, typename Add, typename Sub,
    typename Mul>
void KaratsubaAssignMulImpl(T *dst, const T *lhs, const T *rhs, T *memory, const Add &add, const Sub &sub,
    const Mul &mul)
{
    if constexpr (inSizeLog2 <= thresholdLog2) {
        NaiveAssignMul<addAssign>(dst, lhs, rhs, inSizeLog2, add, sub, mul);
    } else {
        // n is half the length of lhs resp. rhs
        constexpr size_t n = (size_t)1 << (inSizeLog2 - 1);
        // write lhs = a0 + X^n a1 and same for rhs
        T *lower = memory;
        T *rest = memory + (2 * n);

        // set lower to a0 * b0
        KaratsubaAssignMulImpl<inSizeLog2 - 1, thresholdLog2, false>(lower, lhs, rhs, rest, add, sub, mul);
        if constexpr (addAssign) {
            ElementwiseAssign(dst, lower, 2 * n, add);
        } else {
            std::memcpy(dst, lower, 2 * n * sizeof(T));
            std::memset(dst + 2 * n, 0, 2 * n * sizeof(T));
        }
        ElementwiseAssign(dst + n, lower, 2 * n, sub);

        // set upper to a1 * b1
        T *upper = lower; // reuse memory
        KaratsubaAssignMulImpl<inSizeLog2 - 1, thresholdLog2, false>(upper, lhs + n, rhs + n, rest, add, sub, mul);
        ElementwiseAssign(dst + 2 * n, upper, 2 * n, add);
        ElementwiseAssign(dst + n, upper, 2 * n, sub);

        // and now we come to the mid term X^n (a0 * b1 + a1 * b0)
        T *lhsCombined = upper; // reuse memory
        T *rhsCombined = upper + n;
        std::memcpy(lhsCombined, lhs, n * sizeof(T));
        std::memcpy(rhsCombined, rhs, n * sizeof(T));
        ElementwiseAssign(lhsCombined, lhs + n, n, add);
        ElementwiseAssign(rhsCombined, rhs + n, n, add);
        KaratsubaAssignMulImpl<inSizeLog2 - 1, thresholdLog2, true>(dst + n, lhsCombined, rhsCombined, rest, add, sub,
            mul);
    }
}

template <unsigned thresholdLog2, bool addAssign, typename T, typename Add, typename Sub, typename Mul>
void DispatchKaratsubaAssignMulImpl(unsigned inSizeLog2, T *dst, const T *lhs, const T *rhs, T *memory, const Add &add,
    const Sub &sub, const Mul &mul)
{
    if (inSizeLog2 > 16) {
        throw std::invalid_argument("size must be between 1 and 2^16");
    }
    switch (inSizeLog2) {
        case 0:
            KaratsubaAssignMulImpl<0, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub, mul);
            break;
        case 1:
            KaratsubaAssignMulImpl<1, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub, mul);
            break;
        case 2:
            KaratsubaAssignMulImpl<2, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub, mul);
            break;
        case 3:
            KaratsubaAssignMulImpl<3, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub, mul);
            break;
        case 4:
            KaratsubaAssignMulImpl<4, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub, mul);
            break;
        case 5:
            KaratsubaAssignMulImpl<5, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub, mul);
            break;
        case 6:
            KaratsubaAssignMulImpl<6, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub, mul);
            break;
        case 7:
            KaratsubaAssignMulImpl<7, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub, mul);
            break;
        case 8:
            KaratsubaAssignMulImpl<8, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub, mul);
            break;
        case 9:
            KaratsubaAssignMulImpl<9, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub, mul);
            break;
        case 10:
            KaratsubaAssignMulImpl<10, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub,
                mul);
            break;
        case 11:
            KaratsubaAssignMulImpl<11, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub,
                mul);
            break;
        case 12:
            KaratsubaAssignMulImpl<12, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub,
                mul);
            break;
        case 13:
            KaratsubaAssignMulImpl<13, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub,
                mul);
            break;
        case 14:
            KaratsubaAssignMulImpl<14, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub,
                mul);
            break;
        case 15:
            KaratsubaAssignMulImpl<15, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub,
                mul);
            break;
        case 16:
            KaratsubaAssignMulImpl<16, thresholdLog2, addAssign, T, Add, Sub, Mul>(dst, lhs, rhs, memory, add, sub,
                mul);
            break;
    }
}
/*
 * @brief Karatsuba 算法进行多项式乘法
 *
 * 该函数实现了 Karatsuba 算法，通过分治法高效地计算两个多项式或大整数的乘积。与传统的逐元素乘法相比，Karatsuba
 * 算法通过减少乘法次数来提高性能，适用于较大的输入数据。
 *
 * @tparam thresholdLog2 阈值，指定何时切换到传统的乘法算法（较小的输入数据时）。
 * @tparam T 数据类型。
 * @tparam Add 加法操作的类型。
 * @tparam Sub 减法操作的类型。
 * @tparam Mul 乘法操作的类型。
 *
 * @param dst 结果数组，存储计算得到的乘积。
 * @param dstLen 结果数组的长度。
 * @param lhs 左侧多项式或大整数的数组。
 * @param lhsLen 左侧数组的长度。
 * @param rhs 右侧多项式或大整数的数组。
 * @param rhsLen 右侧数组的长度。
 * @param add 加法操作，定义了如何对两个元素进行加法操作。
 * @param sub 减法操作，定义了如何对两个元素进行减法操作。
 * @param mul 乘法操作，定义了如何对两个元素进行乘法操作。
 */
template <unsigned thresholdLog2, typename T, typename Add, typename Sub, typename Mul>
void Karatsuba(T *dst, size_t dstLen, const T *lhs, size_t lhsLen, const T *rhs, size_t rhsLen, const Add &add,
    const Sub &sub, const Mul &mul)
{
    if (lhsLen == 0 || rhsLen == 0) {
        return;
    }
    assert(dstLen >= lhsLen + rhsLen);
    int blockSizeLog2 = (63 - static_cast<int>(std::max(LeadingZeros(lhsLen), LeadingZeros(rhsLen))));
    size_t n = (size_t)1 << blockSizeLog2;
    assert(lhsLen >= n);
    assert(rhsLen >= n);
    assert((lhsLen < 2 * n) || (rhsLen < 2 * n));
    const size_t memElementCount = KaratsubaMemElementCount<T>(static_cast<unsigned int>(blockSizeLog2), thresholdLog2);
    std::unique_ptr<T[]> memory(new T[memElementCount]);

    for (size_t i = 0; i + n <= lhsLen; i += n) {
        for (size_t j = 0; j + n <= rhsLen; j += n) {
            DispatchKaratsubaAssignMulImpl<thresholdLog2, true>(static_cast<unsigned int>(blockSizeLog2), dst + i + j,
                lhs + i, rhs + j, memory.get(), add, sub, mul);
        }
    }

    // now make block sizes smaller and compute the remaining elements
    size_t lhsRem = (lhsLen / n) * n;
    size_t rhsRem = (rhsLen / n) * n;
    blockSizeLog2 -= 1;
    n = n / 2;
    while (blockSizeLog2 >= 0) {
        if (lhsLen >= lhsRem + n) {
            for (size_t j = 0; j + n <= rhsRem; j += n) {
                DispatchKaratsubaAssignMulImpl<thresholdLog2, true>(static_cast<unsigned int>(blockSizeLog2),
                    dst + lhsRem + j, lhs + lhsRem, rhs + j, memory.get(), add, sub, mul);
            }
            lhsRem += n;
        }
        if (rhsLen >= rhsRem + n) {
            for (size_t i = 0; i + n <= lhsLen; i += n) {
                DispatchKaratsubaAssignMulImpl<thresholdLog2, true>(static_cast<unsigned int>(blockSizeLog2),
                    dst + rhsRem + i, lhs + i, rhs + rhsRem, memory.get(), add, sub, mul);
            }
            rhsRem += n;
        }
        n = n / 2;
        blockSizeLog2 -= 1;
    }
}
/*
 * @brief 计算 Karatsuba 算法中所需的内存元素数量
 * 该函数根据输入大小的对数和阈值对数计算 Karatsuba 算法中所需的内存元素数量。通过分治法将问题规模缩小，减少运算次数。
 * @param inSizeLog2 输入数据大小的对数（以 2 为底），即输入数据的位数。
 * @param thresholdLog2 阈值大小的对数（以 2 为底），用于判断何时停止递归。
 * @return 返回 Karatsuba 算法所需的内存元素数量。如果输入大小小于或等于阈值，返回 0。
 */
template <typename T> inline size_t KaratsubaMemElementCount(size_t inSizeLog2, size_t thresholdLog2)
{
    if (inSizeLog2 <= thresholdLog2) {
        return 0;
    }
    const size_t elements = ((size_t)1 << (inSizeLog2 + 1)) - ((size_t)1 << (thresholdLog2 + 1));
    return elements;
}

inline void TestKaratsuba()
{
    auto add = [](int a, int b) { return a + b; };
    auto sub = [](int a, int b) { return a - b; };
    auto mul = [](int a, int b) { return a * b; };
    {
        std::vector<int> lhs = { 0, 1, 1 };
        std::vector<int> rhs = { 1, 0, 0, 1 };
        std::vector<int> dst;
        dst.resize(7);
        Karatsuba<1>(&dst[0], dst.size(), &lhs[0], lhs.size(), &rhs[0], rhs.size(), add, sub, mul);
        std::vector<int> expected = { 0, 1, 1, 0, 1, 1, 0 };
        assert(dst == expected);
    }
    {
        std::vector<int> lhs = { 1, 1, 0, 1, 0, 2, 2, 2 };
        std::vector<int> rhs = { 2, 0, 0, -1, 1, 1, 1 };
        std::vector<int> dst;
        dst.resize(15);
        Karatsuba<1>(&dst[0], dst.size(), &lhs[0], lhs.size(), &rhs[0], rhs.size(), add, sub, mul);
        std::vector<int> expected = { 2, 2, 0, 1, 0, 6, 5, 6, -1, 1, 2, 6, 4, 2, 0 };
        assert(dst == expected);
    }
    {
        std::vector<int> lhs = { 0, 1, 0, 2, 2 };
        std::vector<int> rhs = { 2, 0, 0, -1, 1, 1, 1 };
        std::vector<int> dst;
        dst.resize(12);
        Karatsuba<2>(&dst[0], dst.size(), &lhs[0], lhs.size(), &rhs[0], rhs.size(), add, sub, mul);
        std::vector<int> expected = { 0, 2, 0, 4, 3, 1, -1, 1, 4, 4, 2, 0 };
        assert(dst == expected);
    }
    std::cout << "TestKaratsuba(): success" << std::endl;
}
}
#endif